package com.boarsoft.dao.sharding.jdbc;

import java.io.Closeable;
import java.io.IOException;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;

import javax.sql.DataSource;

import org.springframework.util.Assert;

public class ConnectionExPool implements Closeable {
	protected DataSource dataSource;
	protected final BlockingQueue<ConnectionEx> idleQueue = new LinkedBlockingQueue<ConnectionEx>();
	protected final ThreadLocal<ConnectionEx> holder = new ThreadLocal<ConnectionEx>();
	protected int validTimeout = 1;

	public ConnectionEx getConnection() throws SQLException {
		ConnectionEx connEx = holder.get();
		if (connEx == null) {
			Connection conn = dataSource.getConnection();
			Assert.notNull(conn, "Connection can not be null");
			Assert.isTrue(conn.isValid(validTimeout), "Connection is invalid");
			connEx = new ConnectionEx(conn);
			holder.set(connEx);
		}
		return connEx;
	}

	public void returnConnection(ConnectionEx connEx) throws SQLException {
		idleQueue.add(connEx);
	}

	@Override
	public void close() throws IOException {
		for (ConnectionEx connEx : idleQueue) {
			connEx.close(true);
		}
		idleQueue.clear();
		holder.remove();
	}
}
